package org.znxs.znmanus.agent;

import cn.hutool.core.util.StrUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import org.znxs.znmanus.model.ZNAgentStatus;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;


/**
 * 基础智能体类
 * 作用是定义智能体的属性 以及管理代理状态和执行流程
 * <p>
 * 提供状态转换、内存管理和基于步骤的执行循环的基础功能。
 * 子类必须实现step方法。
 */
@Data
@Slf4j
public abstract class ZNBaseAgent {

    /**
     * 智能体名称
     */
    private String agentName;

    /**
     * 系统预设（最开始的定义，即该模型是干嘛的）
     */
    private String systemPrompt;

    /**
     * 执行下一步的时候的提示（即提示模型该怎么一步一步执行）
     */
    private String nextPrompt;

    /**
     * 状态 默认空闲状态
     */
    private ZNAgentStatus status = ZNAgentStatus.IDLE;

    /**
     * 当前执行步数
     */
    private Integer currentStep = 0;

    /**
     * 最大步数
     */
    private Integer maxSteps = 10;


    /**
     * 调用的大模型 LLM
     */
    private ChatClient chatClient;

    /**
     * 自主维护上下文
     */
    private List<Message> messages = new ArrayList<Message>();

    /**
     * 运行代理
     *
     * @param userPrompt 用户提示词
     * @return 展示最后的结果
     */
    public String run(String userPrompt) {
        // 状态判断
        if (this.status != ZNAgentStatus.IDLE) {
            throw new RuntimeException(String.format("智能体不为空闲状态，当前状态%s，请稍后再试", this.status));
        }
        // 基础参数校验
        if (StrUtil.isBlank(userPrompt)) {
            throw new RuntimeException("用户提示词为空，请输入用户提示词");
        }
        // 更改状态
        this.status = ZNAgentStatus.RUNNING;
        // 初始化最初的用户提示词到消息上下文中
        this.messages.add(new UserMessage(userPrompt));

        // 初始化保存执行返回结果的集合
        List<String> results = new ArrayList<>();

        try {
            // 代理调用 进行循环调用
            for (int i = 0; i < maxSteps && status != ZNAgentStatus.FINISHED; i++) {
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("执行步骤：[{} / {}]", stepNumber, maxSteps);
                // 调用单步执行  并将执行结果拼接到返回集合中
                String stepResult = this.step();
                String result = String.format("Step: %s, Result: %s", stepNumber, stepResult);
                results.add(result);
            }

            // 检查是否超出最大步数限制
            if (currentStep >= maxSteps) {
                this.status = ZNAgentStatus.FINISHED;
                results.add(String.format("终止：到达最大步数[%s / %s]", this.currentStep, this.maxSteps));
            }

            // 插入拼接的字符串 返回处理后的 结果
            return String.join("\n", results);
        } catch (Exception e) {
            status = ZNAgentStatus.ERROR;
            log.error("智能体生成出现错误", e);
            return "智能体[BASE]执行出现错误" + e.getMessage();
        } finally {
            // 清理资源
            this.cleanUp();
        }

    }

    /**
     * 运行代理（流失处理）
     *
     * @param userPrompt 用户提示词
     * @return 展示最后的结果
     */
    public SseEmitter runStream(String userPrompt) {
        // 创建一个超时时间合适的SseEmitter
        SseEmitter sseEmitter = new SseEmitter(300000L);

        // 这里需要使用异步线程进行处理，因为需要一步一步进行返回，所以需要异步处理
        // 使用的线程异步
        CompletableFuture.runAsync(() -> {
            try {
                // 状态判断
                if (this.status != ZNAgentStatus.IDLE) {
                    // 错误校验使用SseEmitter 进行处理
                    sseEmitter.send(String.format("智能体不为空闲状态，当前状态%s，请稍后再试", this.status));
                    sseEmitter.complete();
                    // 注意，这里的return 是回调方法的return，并不是整个方法的return
                    return;
                }
                // 基础参数校验
                if (StrUtil.isBlank(userPrompt)) {
                    sseEmitter.send("用户提示词为空，请输入用户提示词");
                    sseEmitter.complete();
                    return;
                }
            } catch (IOException e) {
                sseEmitter.completeWithError(e);
            }
            // 更改状态
            this.status = ZNAgentStatus.RUNNING;
            // 初始化最初的用户提示词到消息上下文中
            this.messages.add(new UserMessage(userPrompt));


            try {
                // 代理调用 进行循环调用
                for (int i = 0; i < maxSteps && status != ZNAgentStatus.FINISHED; i++) {
                    int stepNumber = i + 1;
                    currentStep = stepNumber;
                    log.info("执行步骤：[{} / {}]", stepNumber, maxSteps);
                    // 调用单步执行  直接将结果使用SseEmitter发射出去
                    String stepResult = this.step();
                    String result = String.format("Step: %s, Result: %s", stepNumber, stepResult);
                    sseEmitter.send(result);
                }
                // 检查是否超出最大步数限制
                if (currentStep >= maxSteps) {
                    this.status = ZNAgentStatus.FINISHED;
                    sseEmitter.send(String.format("终止：到达最大步数[%s / %s]", this.currentStep, this.maxSteps));
                }
                // 为什么在这里调用完成方法？ 因为complete的作用是告诉服务端我已经发送完所有数据了，避免一直挂起，占用资源
                sseEmitter.complete();

            } catch (Exception e) {
                status = ZNAgentStatus.ERROR;
                log.error("智能体生成出现错误", e);
                try {
                    sseEmitter.send("用户提示词为空，请输入用户提示词");
                    sseEmitter.complete();
                } catch (IOException ex) {
                    sseEmitter.completeWithError(ex);
                }
            } finally {
                // 清理资源
                this.cleanUp();
            }
        });

        // 当超过规定时间没有发送 就会调用
        sseEmitter.onTimeout(() -> {
            this.status = ZNAgentStatus.IDLE;
            this.cleanUp();
            log.warn("SSE 连接超时");
        });

        // 当完成的时候就会调用
        sseEmitter.onCompletion(() -> {
            if (this.status == ZNAgentStatus.RUNNING) {
                this.status = ZNAgentStatus.FINISHED;
            }
            this.cleanUp();
            log.info("SSE 连接完成");
        });

        return sseEmitter;
    }

    /**
     * 执行单个步骤
     *
     * @return 执行单步结果
     */
    public abstract String step();


    /**
     * 资源清理
     */
    protected void cleanUp() {
        // 子类可以实现此方法来清理资源
    }


}
